use std::env;
use std::time::Instant;

use dora_node_api::dora_core::config::DataId;
use dora_node_api::{DoraNode, Event, IntoArrow};
use eyre::Context;
use image::imageops::FilterType;
use image::GenericImageView;
use ort::{GraphOptimizationLevel, Session};

use runner::{convert_result, pixels_to_array};

fn main() -> eyre::Result<()> {
    let data_id = DataId::from("boxes".to_owned());
    let (mut node, mut events) = DoraNode::init_from_env()?;

    let model_path = env::var("model")?;
    let names = env::var("names")?;
    let names: Vec<&str> = names.trim().split(",").map(|v| v.trim()).collect();

    let model = Session::builder()?
        .with_optimization_level(GraphOptimizationLevel::Level3)?
        .with_intra_threads(4)?
        .commit_from_file(model_path)?;

    while let Some(event) = events.recv() {
        match event {
            Event::Input { id, metadata, data } => match id.as_str() {
                "image" => {
                    let data: &[u8] = TryFrom::try_from(&data).context("expected &[u8] message")?;

                    let now = Instant::now();
                    let original_img = image::load_from_memory(data)?;
                    println!("load use time: {}", now.elapsed().as_millis());
                    let resize_width = 640;
                    let resize_height = 640;
                    let now = Instant::now();
                    let image = original_img.resize_exact(
                        resize_width,
                        resize_height,
                        FilterType::Nearest,
                    );
                    println!("resize use time: {}", now.elapsed().as_millis());

                    let now = Instant::now();
                    let input = pixels_to_array(image);
                    println!("covert use time: {}", now.elapsed().as_millis());

                    let now = Instant::now();
                    let outputs = model.run(ort::inputs!["images" => input.view()]?)?;
                    println!("model use time: {}", now.elapsed().as_millis());

                    let result = convert_result(&outputs, &names, &original_img)?;
                    println!("result: {:?}", result);

                    node.send_output(
                        data_id.clone(),
                        metadata.parameters,
                        "".into_arrow(),
                    )?;
                }
                other => eprintln!("Ignoring unexpected input `{other}`"),
            },
            Event::Stop => {
                println!("Received manual stop");
            }
            Event::InputClosed { id } => {
                println!("Input `{id}` was closed");
            }
            other => eprintln!("Received unexpected input: {other:?}"),
        }
    }

    Ok(())
}
